import org.bytedeco.javacv.*;
import java.io.*;
import java.net.*;
import java.nio.ByteBuffer;
import java.util.*;
import java.util.concurrent.*;
import java.util.zip.CRC32;

/**
 * RTSP流转交通部标JT/T 1078协议转发器
 * 实现RTSP流抓取并按照JT/T 1078协议转发到监管平台
 */
public class RtspToJtt1078Forwarder {
    // 配置参数
    private final String rtspUrl;
    private final String platformIp;
    private final int platformPort;
    private final String terminalPhone; // 终端手机号(设备唯一标识)
    private final int channelId; // 通道ID
    private final String localIp;
    private final int localPort;

    // 协议常量
    private static final byte MSG_TYPE_VIDEO = 0x01; // 视频消息
    private static final byte ENCODING_H264 = 0x01; // H.264编码
    private static final byte STREAM_TYPE_MAIN = 0x01; // 主码流
    private static final short PACKAGE_MAX_SIZE = 1024; // 分包最大尺寸

    // 运行状态
    private volatile boolean isRunning = false;
    private DatagramSocket udpSocket;
    private ScheduledExecutorService executor;
    private FrameGrabber grabber;
    private int frameSequence = 0; // 帧序列号

    public RtspToJtt1078Forwarder(String rtspUrl, String platformIp, int platformPort,
                                 String terminalPhone, int channelId, String localIp, int localPort) {
        this.rtspUrl = rtspUrl;
        this.platformIp = platformIp;
        this.platformPort = platformPort;
        this.terminalPhone = terminalPhone;
        this.channelId = channelId;
        this.localIp = localIp;
        this.localPort = localPort;
    }

    /**
     * 启动转发器
     */
    public void start() throws IOException {
        isRunning = true;
        udpSocket = new DatagramSocket(localPort, InetAddress.getByName(localIp));
        executor = Executors.newScheduledThreadPool(3);
        
        // 启动RTSP拉流和处理
        executor.execute(this::processRtspStream);
        
        // 启动心跳发送(每30秒一次)
        executor.scheduleAtFixedRate(this::sendHeartbeat, 0, 30, TimeUnit.SECONDS);
        
        System.out.println("JT/T 1078转发器已启动，终端号: " + terminalPhone);
    }

    /**
     * 处理RTSP流并转发
     */
    private void processRtspStream() {
        try {
            // 初始化RTSP抓取器
            grabber = FrameGrabber.createDefault(rtspUrl);
            grabber.setOption("rtsp_transport", "tcp");
            grabber.setOption("stimeout", "5000000"); // 5秒超时
            grabber.start();
            
            System.out.println("RTSP流拉取成功: " + rtspUrl);
            System.out.println("视频信息 - 宽: " + grabber.getImageWidth() + 
                               ", 高: " + grabber.getImageHeight() + 
                               ", 帧率: " + grabber.getFrameRate());

            // 发送终端注册消息
            sendRegister();

            // 处理视频帧
            Frame frame;
            while (isRunning && (frame = grabber.grab()) != null) {
                // 转换为H.264编码(JT/T 1078主要支持H.264)
                byte[] h264Data = convertToH264(frame);
                if (h264Data != null && h264Data.length > 0) {
                    // 发送H.264数据到平台
                    sendVideoData(h264Data, System.currentTimeMillis());
                    frameSequence++;
                }
                
                // 控制发送速率
                Thread.sleep(40); // 约25fps，可根据需要调整
            }
        } catch (Exception e) {
            System.err.println("RTSP流处理错误: " + e.getMessage());
            e.printStackTrace();
        } finally {
            stop();
        }
    }

    /**
     * 转换帧数据为H.264格式
     */
    private byte[] convertToH264(Frame frame) {
        // 实际应用中需要使用编码器将帧数据转换为H.264
        // 这里使用JavaCV的编码器进行转换
        try (FrameConverter<Mat> converter = new OpenCVFrameConverter.ToMat();
             FFmpegFrameRecorder recorder = new FFmpegFrameRecorder(
                 new NullOutputStream(), 
                 frame.imageWidth, 
                 frame.imageHeight, 
                 0)) {
            
            recorder.setVideoCodec(avcodec.AV_CODEC_ID_H264);
            recorder.setFormat("h264");
            recorder.setFrameRate(25);
            recorder.setVideoBitrate(2000000); // 2Mbps
            recorder.start();
            
            // 使用内存输出流捕获编码后的H.264数据
            ByteArrayOutputStream out = new ByteArrayOutputStream();
            recorder.setOutputStream(new DataOutputStream(out));
            recorder.record(frame);
            recorder.stop();
            
            return out.toByteArray();
        } catch (Exception e) {
            System.err.println("H.264编码错误: " + e.getMessage());
            return null;
        }
    }

    /**
     * 发送视频数据到平台
     */
    private void sendVideoData(byte[] h264Data, long timestamp) throws IOException {
        // 计算需要分多少包
        int totalPackets = (int) Math.ceil((double) h264Data.length / PACKAGE_MAX_SIZE);
        int remaining = h264Data.length;
        
        // 生成当前帧的唯一标识
        long frameId = System.currentTimeMillis();
        
        for (int i = 0; i < totalPackets; i++) {
            int offset = i * PACKAGE_MAX_SIZE;
            int length = Math.min(remaining, PACKAGE_MAX_SIZE);
            remaining -= length;
            
            // 创建数据包
            byte[] packet = createVideoPacket(
                h264Data, offset, length, i, totalPackets, frameId, timestamp);
            
            // 发送UDP包
            DatagramPacket datagram = new DatagramPacket(
                packet, packet.length, 
                InetAddress.getByName(platformIp), platformPort);
            udpSocket.send(datagram);
        }
    }

    /**
     * 创建JT/T 1078协议的视频数据包
     */
    private byte[] createVideoPacket(byte[] data, int offset, int length, 
                                    int packetIndex, int totalPackets, 
                                    long frameId, long timestamp) {
        try (ByteArrayOutputStream out = new ByteArrayOutputStream();
             DataOutputStream dos = new DataOutputStream(out)) {
            
            // 1. 消息头 (28字节)
            dos.writeShort(0x3031); // 起始标识 0x3031
            dos.writeShort(28 + length); // 消息体总长度
            
            // 消息ID (0x9101表示实时视频)
            dos.writeShort(0x9101);
            
            // 终端手机号(设备ID)
            byte[] phoneBytes = terminalPhone.getBytes("GBK");
            byte[] phoneField = new byte[10];
            System.arraycopy(phoneBytes, 0, phoneField, 0, Math.min(phoneBytes.length, 10));
            dos.write(phoneField);
            
            // 通道ID
            dos.writeByte(channelId);
            
            // 消息体属性
            short property = (short) (length << 10); // 数据长度
            property |= (ENCODING_H264 << 7); // 编码方式
            property |= (STREAM_TYPE_MAIN << 5); // 码流类型
            property |= (totalPackets > 1 ? 1 : 0); // 是否分包
            dos.writeShort(property);
            
            // 2. 分包信息 (如需要)
            if (totalPackets > 1) {
                dos.writeInt((int) frameId); // 帧编号
                dos.writeByte(totalPackets); // 总包数
                dos.writeByte(packetIndex + 1); // 当前包序号
            }
            
            // 3. 时间戳
            dos.writeLong(timestamp);
            
            // 4. 视频数据
            dos.write(data, offset, length);
            
            // 5. 校验码 (CRC32)
            CRC32 crc = new CRC32();
            crc.update(out.toByteArray());
            dos.writeInt((int) crc.getValue());
            
            return out.toByteArray();
        } catch (IOException e) {
            throw new RuntimeException("创建视频数据包失败", e);
        }
    }

    /**
     * 发送终端注册消息
     */
    private void sendRegister() throws IOException {
        try (ByteArrayOutputStream out = new ByteArrayOutputStream();
             DataOutputStream dos = new DataOutputStream(out)) {
            
            // 消息头
            dos.writeShort(0x3031); // 起始标识
            dos.writeShort(38); // 消息体总长度(28+10)
            dos.writeShort(0x0100); // 消息ID(终端注册)
            
            // 终端手机号
            byte[] phoneBytes = terminalPhone.getBytes("GBK");
            byte[] phoneField = new byte[10];
            System.arraycopy(phoneBytes, 0, phoneField, 0, Math.min(phoneBytes.length, 10));
            dos.write(phoneField);
            
            // 通道ID(0表示终端)
            dos.writeByte(0);
            
            // 消息体属性(无分包)
            dos.writeShort(0);
            
            // 时间戳
            dos.writeLong(System.currentTimeMillis());
            
            // 消息体(终端注册信息)
            dos.write("JT1078转发器".getBytes("GBK")); // 终端制造商(10字节)
            
            // 校验码
            CRC32 crc = new CRC32();
            crc.update(out.toByteArray());
            dos.writeInt((int) crc.getValue());
            
            // 发送注册消息
            byte[] data = out.toByteArray();
            DatagramPacket packet = new DatagramPacket(
                data, data.length, 
                InetAddress.getByName(platformIp), platformPort);
            udpSocket.send(packet);
            
            System.out.println("已发送终端注册消息");
        }
    }

    /**
     * 发送心跳消息
     */
    private void sendHeartbeat() {
        try (ByteArrayOutputStream out = new ByteArrayOutputStream();
             DataOutputStream dos = new DataOutputStream(out)) {
            
            // 消息头
            dos.writeShort(0x3031); // 起始标识
            dos.writeShort(28); // 消息体总长度(无消息体)
            dos.writeShort(0x0001); // 消息ID(心跳)
            
            // 终端手机号
            byte[] phoneBytes = terminalPhone.getBytes("GBK");
            byte[] phoneField = new byte[10];
            System.arraycopy(phoneBytes, 0, phoneField, 0, Math.min(phoneBytes.length, 10));
            dos.write(phoneField);
            
            // 通道ID(0表示终端)
            dos.writeByte(0);
            
            // 消息体属性(无分包)
            dos.writeShort(0);
            
            // 时间戳
            dos.writeLong(System.currentTimeMillis());
            
            // 校验码
            CRC32 crc = new CRC32();
            crc.update(out.toByteArray());
            dos.writeInt((int) crc.getValue());
            
            // 发送心跳消息
            byte[] data = out.toByteArray();
            DatagramPacket packet = new DatagramPacket(
                data, data.length, 
                InetAddress.getByName(platformIp), platformPort);
            udpSocket.send(packet);
            
        } catch (Exception e) {
            System.err.println("发送心跳消息失败: " + e.getMessage());
        }
    }

    /**
     * 停止转发器
     */
    public void stop() {
        isRunning = false;
        
        if (executor != null) {
            executor.shutdownNow();
        }
        
        if (grabber != null) {
            try {
                grabber.stop();
                grabber.release();
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
        
        if (udpSocket != null && !udpSocket.isClosed()) {
            udpSocket.close();
        }
        
        System.out.println("JT/T 1078转发器已停止");
    }

    public static void main(String[] args) {
        try {
            // 配置参数
            String rtspUrl = "rtsp://example.com/stream"; // 替换为实际RTSP地址
            String platformIp = "192.168.1.10"; // 交通部标平台IP
            int platformPort = 5000; // 交通部标平台端口
            String terminalPhone = "13800138000"; // 终端手机号(设备标识)
            int channelId = 1; // 通道ID
            String localIp = "192.168.1.100"; // 本地IP
            int localPort = 6000; // 本地端口

            // 创建并启动转发器
            RtspToJtt1078Forwarder forwarder = new RtspToJtt1078Forwarder(
                rtspUrl, platformIp, platformPort, 
                terminalPhone, channelId, localIp, localPort);
            
            forwarder.start();
            
            // 添加关闭钩子
            Runtime.getRuntime().addShutdownHook(new Thread(forwarder::stop));
            
        } catch (Exception e) {
            System.err.println("转发器启动失败: " + e.getMessage());
            e.printStackTrace();
        }
    }

    /**
     * 空输出流，用于捕获编码器输出
     */
    private static class NullOutputStream extends OutputStream {
        @Override
        public void write(int b) {}
    }
}
